import gradio as gr
import numpy as np
import random
import torch
from diffusers import StableDiffusion3Pipeline, SD3Transformer2DModel, FlowMatchEulerDiscreteScheduler
import spaces
import qwen_api

device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16
repo = "hf-models/stable-diffusion-3-medium-diffusers"


pipe = StableDiffusion3Pipeline.from_pretrained(
    repo, torch_dtype=dtype).to(device)
MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 1344


@spaces.GPU
def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, use_zh, progress=gr.Progress(track_tqdm=True)):
    print("使用中文", use_zh)
    resPrompt = ""
    if (use_zh):
        try:
            qwenRes = qwen_api.query({
                "messages": [
                    {
                        "role": "system",
                        "content": "你是一个文生图 Prompt 翻译器，将中文翻译为英文，翻译简洁精准，如果已经是英文就直接原文返回。你只需要直接给出英文答案, 不要废话"
                    },
                    {
                        "role": "user",
                        "content": prompt
                    }
                ],
                "stream": False

            })
            resPrompt = qwenRes["choices"][0]["message"]["content"]
        except Exception as e:
            print("请求失败")
            print(f"Error: {e}")
            resPrompt = prompt
            raise gr.Error("解析中文提示词失败，请稍后或取消勾选中文提示词")
    else:
        resPrompt = prompt

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    print(prompt)
    print(resPrompt)

    generator = torch.Generator().manual_seed(seed)

    image = pipe(
        prompt=resPrompt,
        negative_prompt=negative_prompt,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        width=width,
        height=height,
        generator=generator
    ).images[0]

    torch.cuda.empty_cache()

    return image, seed


examples = [
    "美丽的白色樱花盛开，在晴朗的天空背景下，树枝上精致而新鲜的花朵，柔和对焦的摄影，日式风格，微距镜头，自然光，春天的气息，色调清新",
    "杰作、获奖、专业、高度详细、居中、中景、全身、动漫风格、插图、风格化、油画、女巫女孩、可爱、拥抱、抱着月亮、太空行星背景、魔法阵列、分形艺术、超级极繁主义、（史诗构图、史诗比例、超现实主义）、鲜艳的色彩、自然光、景深、高清、64K、全景",
    "美食摄影，沙卡蔬卡，顶视图，豪华的米其林厨房风格，工作室灯光，景深，超详细",    "苹果树下一只美丽的老虎神奇宝贝的数字艺术，卡通风格，哑光绘画，魔幻现实主义，鲜艳色彩，超高品质，高细节，高分辨率",
    "一只鹰在森林中飞翔，高分辨率，超现实主义的照片，采用专业色彩分级风格，具有柔和的阴影和干净，清晰的焦点数码摄影, 8k",
    "一张杂志照片，尼康 500，拍摄于日出时的瀑布",
    "小丑是蝙蝠侠电影系列中的反派，他在一面破碎的镜子中看着自己，风格类似克里斯托弗·诺兰，充满动态、黑暗和恐怖的氛围，戏剧性的灯光，细节丰富，充满隐藏细节，超现实主义，壁纸肖像，色彩缤纷的现实主义，对细节的现实关注，高度细节化，色彩缤纷的现实主义，超高清，8K",
    "一盘肉桂卷，旁边放着一些糖，放在一张古老的蓝色桌面上，背景是红布，顶视图，食物摄影，电影风格的照片，高分辨率。",
    "1女孩，户外，独奏，花，蓝眼睛，长发，看着观众，长裙，帽子，海洋，连衣裙，坐着，白花，天空，虫子，帽子花，白天，刘海，蝴蝶，水，长袖，蓝天，岩石，云，橙色连衣裙，金发，地平线，橙色头饰，头饰，闭嘴，微笑，头花，短袖，看向一边，鸟，风景，多山的地平线，红色头饰，白色连衣裙",
    "Some lychees are on the plate, crystal clear, beautiful, appetizing, shiny",
    "Crystal Cat",
    "1girl, outdoors, solo, flower, blue eyes, long hair, looking at viewer, long dress，hat, ocean, dress, sitting, white flower, sky, bug, hat flower, day, bangs, butterfly, water, long sleeves, blue sky, rock, cloud, orange dress, blonde hair, horizon, orange headwear, hair ornament, closed mouth, smile, hair flower, short sleeves, looking to the side, bird, scenery, mountainous horizon, red headwear, white dress",
    "Long exposure photo of Tokyo street, blurred motion, streaks of light, surreal, dreamy, ghosting effect, highly detailed.",
    "A living room, bright modern Scandinavian style house, large windows.",
    "A dog dressed in a prisoner's uniform, standing in front of a row of jail cells with a sign on the door saying 'Thief Dog', suggesting a humorous take on canine mischief.",
    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
    "An astronaut riding a green horse",
    "A delicious ceviche cheesecake slice",
]
css = """
#col-container {
    margin: 0 auto;
    max-width: 580px;
}
"""
with gr.Blocks(css=css) as demo:

    with gr.Column(elem_id="col-container"):
        gr.Markdown(f"""
        # Gitee AI Stable Diffusion 3 Medium  
        - 支持使用中文提示词, 由 [Gitee AI Serverless API](https://ai.gitee.com/serverless-api) 支持
        - [前往使用文生图 Prompt 辅助工具](https://ai.gitee.com/apps/stringify/image-prompt-tool)
        """)

        with gr.Row():

            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                lines=3,
                placeholder="输入你的提示词",
                container=False,
            )

            run_button = gr.Button("运行", scale=0)

        with gr.Row():
            use_zh = gr.Checkbox(label="使用中文提示词", value=True)

        result = gr.Image(label="Result", show_label=False)
        with gr.Accordion("环境配置", open=False):

            negative_prompt = gr.Text(
                label="否定提示词",
                value="Disabled feet, abnormal feet, deformed, distorted, disfigured, poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, mutated hands and fingers, disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation, NSFW",
                max_lines=1,
                placeholder="输入否定提示词",
            )

            seed = gr.Slider(
                label="种子",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=0,
            )

            randomize_seed = gr.Checkbox(label="使用随机种子", value=True)

            with gr.Row():

                width = gr.Slider(
                    label="宽度",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=64,
                    value=1024,
                )

                height = gr.Slider(
                    label="高度",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=64,
                    value=1024,
                )

            with gr.Row():

                guidance_scale = gr.Slider(
                    label="引导比例",
                    minimum=0.0,
                    maximum=10.0,
                    step=0.1,
                    value=7.0,
                )

                num_inference_steps = gr.Slider(
                    label="推理步骤数",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=28,
                )

        gr.Examples(
            label="例子",
            examples=examples,
            inputs=[prompt]
        )
    gr.on(
        triggers=[run_button.click, prompt.submit, negative_prompt.submit],
        fn=infer,
        inputs=[prompt, negative_prompt, seed, randomize_seed,
                width, height, guidance_scale, num_inference_steps, use_zh],
        outputs=[result, seed]
    )

demo.queue("auto")
demo.launch(show_api=False)
